#include "TextureProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "LocalEntity.h"
#include "Render/VulkanTool.h"

void TextureProcess::OnInit()
{
	mLocalEntitys = mGameScene->GetEntitysByType<LocalEntity, Model>();

	UboBase ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.view = mGameScene->GetCamera().GetView();

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, ubo](VkCommandBuffer commandBuffer)
		{

			for (uint32 i = 0; i < mLocalEntitys.size(); ++i)
			{
				Model& model = mLocalEntitys[i]->GetComponent<Model>();

				std::vector<VertexPosTexCoord> tempVertexs;
				tempVertexs.resize(model.mVertexs.size());
				for (uint32 i = 0; i < model.mVertexs.size(); ++i)
				{
					tempVertexs[i].pos = model.mVertexs[i].Position;
					tempVertexs[i].texCoord = model.mVertexs[i].TexC;
				}

				mLocalEntitys[i]->mVertexBuffer = VertexBuffer::Builder()
					.SetData(tempVertexs.data())
					.SetSize(tempVertexs.size() * sizeof(VertexPosTexCoord))
					.Build();

				mLocalEntitys[i]->mIndexBuffer = IndexBuffer::Builder()
					.SetData(model.mIndices.data())
					.SetSize(model.mIndices.size() * sizeof(uint32))
					.Build();

				mLocalEntitys[i]->mIndiceSize = model.mIndices.size();

				UboBase entityUbo = ubo;
				entityUbo.model = mLocalEntitys[i]->GetTransform().GetLocalMatrix();

				mLocalEntitys[i]->mUboBuffer = CustomBuffer::Builder()
					.SetData(&entityUbo)
					.SetSize(sizeof(UboBase))
					.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
					.SetMultiple(true)
					.Build();
			}

			mLeftTexture = GetTexture("ComputeProcess", "leftTexture");

			mRightTexture = GetTexture("ComputeProcess", "rightTexture");

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<TextureVertexShader>()
				.PushShader<TextureFragmentShader>()
				.SetDescriptorCount(0, 2)
				.Build();

			mSampler = RHISampler::Builder()
				.SetAddressModeU(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetAddressModeV(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetAddressModeW(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetMaxAnisotropy(1.f)
				.SetMaxLod(1)
				.SetBorderColor(VK_BORDER_COLOR_FLOAT_OPAQUE_WHITE)
				.Build();

			mDescriptor->SetValue("ubo", mLocalEntitys[0]->mUboBuffer);
			mDescriptor->SetValue("samplerColor", mLeftTexture, mSampler);
			mLocalEntitys[0]->mDescriptorMask = mDescriptor->AllocateDescriptor();

			mDescriptor->SetValue("ubo", mLocalEntitys[1]->mUboBuffer);
			mDescriptor->SetValue("samplerColor", mRightTexture, mSampler);
			mLocalEntitys[1]->mDescriptorMask = mDescriptor->AllocateDescriptor();


			mColorRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetMultiple(true)
				.Build(this, "ColorRenderTarget");

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.SetMultiple(true)
				.Build();

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mColorRenderTarget)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mColorRenderTarget)
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void TextureProcess::OnTick(float deltaTime)
{
	std::array<UboBase, 2> ubos;
	for (uint32 i = 0; i < 2; ++i)
	{
		ubos[i].proj = mGameScene->GetCamera().GetProj();
		ubos[i].view = mGameScene->GetCamera().GetView();
		ubos[i].model = mLocalEntitys[i]->GetTransform().GetLocalMatrix();
	}

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubos](VkCommandBuffer commandBuffer)
		{
			std::array<UboBase, 2> entityUbos;
			entityUbos[0] = ubos[0];
			entityUbos[1] = ubos[1];

			for (uint32 i = 0; i < 2; ++i)
			{
				mLocalEntitys[i]->mUboBuffer->WriteMemory(&entityUbos[i]);
			}

#if 0

			VkImageMemoryBarrier imageMemoryBarrier{};
			imageMemoryBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
			imageMemoryBarrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
			imageMemoryBarrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
			imageMemoryBarrier.oldLayout = VK_IMAGE_LAYOUT_GENERAL;
			imageMemoryBarrier.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
			imageMemoryBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
			imageMemoryBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
			imageMemoryBarrier.image = mLeftTexture->GetImage();
			imageMemoryBarrier.subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT , 0, 1, 0, 1 };

			vkCmdPipelineBarrier(
				commandBuffer,
				VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
				VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
				0,
				0, nullptr,
				0, nullptr,
				1, &imageMemoryBarrier
			);

			imageMemoryBarrier.image = mRightTexture->GetImage();

			vkCmdPipelineBarrier(
				commandBuffer,
				VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
				VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
				0,
				0, nullptr,
				0, nullptr,
				1, &imageMemoryBarrier
			);

#else

			TextureTransfer textureTransfer;
			textureTransfer
				.SetSrcStageMask(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_WRITE_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetNewLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetTexture(mLeftTexture)
				.Transfer(commandBuffer);

			textureTransfer
				.SetTexture(mRightTexture)
				.Transfer(commandBuffer);
#endif


			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			for (const auto& entity : mLocalEntitys)
			{
				vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mDescriptor->GetPipelineLayout(), 0, 1, &entity->mDescriptorMask.GetDescriptorSet(), 0, nullptr);

				VkBuffer vertexBuffer = entity->mVertexBuffer->GetBuffer();
				VkDeviceSize offsets = 0;
				vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

				VkBuffer indexBuffer = entity->mIndexBuffer->GetBuffer();
				vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

				vkCmdDrawIndexed(commandBuffer, entity->mIndiceSize, 1, 0, 0, 0);
			}

			vkCmdEndRenderPass(commandBuffer);

		});
}

void TextureProcess::OnResize()
{
	mColorRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.SetMultiple(true)
		.Build(this, "ColorRenderTarget");

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.SetMultiple(true)
		.Build();

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mColorRenderTarget)
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}
